add sliding window support for webgpu gqa by guschmue · Pull Request #25372 · microsoft/onnxruntime

guschmue · 2025-07-11T19:29:47Z

No description provided.

Copilot

Pull Request Overview

Adds support for a sliding window mask in the GPU-based group query attention (GQA) softmax shader, allowing attention to be constrained to a local window of past tokens.

Propagate a new local_window_size parameter through ComputeInternal, ApplyAttention, and ComputeInPlaceSoftmax.
Extend InPlaceSoftmaxProgram to accept local_window_size and emit conditional shader code for sliding-window masking.
Update function signatures and uniform variable setup to include the new window size parameter.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
contrib_ops/webgpu/bert/group_query_attention.cc	Added `local_window_size_` guard for FlashAttention and pass it to ApplyAttention.
contrib_ops/webgpu/bert/attention_common.h	Extended `ApplyAttention` signature to accept `local_window_size`.
contrib_ops/webgpu/bert/attention.h	Updated `InPlaceSoftmaxProgram` constructor and added `local_window_size_` field.
contrib_ops/webgpu/bert/attention.cc	Injected sliding-window logic into shader code and passed `local_window_size` through to the program.

Comments suppressed due to low confidence (2)

onnxruntime/contrib_ops/webgpu/bert/attention.cc:228

[nitpick] Consider adding unit or integration tests to cover the new sliding window behavior in InPlaceSoftmaxProgram to ensure correctness under various window sizes, including edge cases like window size zero or larger than sequence length.

  bool has_sliding_window = local_window_size_ != -1;

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc:199

[nitpick] Align the indentation for this condition with the other conditions in the if statement to improve readability (e.g., match the two-space indent used for head_sink == nullptr).

      local_window_size_ == -1 &&

onnxruntime/contrib_ops/webgpu/bert/attention.cc

Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>

snnn · 2025-07-25T16:36:21Z

Hi there! We haven't cut the release branch for this version yet, so I'm removing the release:1.23.0 label for now to keep things tidy. Thanks so much for your contribution! We'll make sure this gets included when the release is prepared. 🤖

…gth (#25594) ### Description  #25372 adds sliding window support for Group Query Attention, disabling Flash Attention as it's not yet supported. This PR adds a check for the sliding window and applies Flash Attention when the window size exceeds the KV cache length or total sequence length. ### Motivation and Context See above.

Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>

…gth (microsoft#25594) ### Description  microsoft#25372 adds sliding window support for Group Query Attention, disabling Flash Attention as it's not yet supported. This PR adds a check for the sliding window and applies Flash Attention when the window size exceeds the KV cache length or total sequence length. ### Motivation and Context See above.

add sliding window support for webgpu gqa

adfe338

guschmue added the ep:WebGPU ort-web webgpu provider label Jul 11, 2025

guschmue requested a review from Copilot July 11, 2025 19:29

Copilot AI reviewed Jul 11, 2025

View reviewed changes

onnxruntime/contrib_ops/webgpu/bert/attention.cc Show resolved Hide resolved

onnxruntime/contrib_ops/webgpu/bert/attention.cc Show resolved Hide resolved

fs-eire reviewed Jul 14, 2025

View reviewed changes